/*
 * Memory management routines
 * 24 Mar 2012
 * Frio
 */

#include "MemoryManager.h"
#include <stdio.h>
#include <malloc.h>
#include <unistd.h>
#include <signal.h>
#include <stdio.h>
#include <malloc.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/mman.h>

using namespace std;

#define handle_error(msg) do { perror(msg); exit(EXIT_FAILURE); } while (0)

// Prototypes
static void *customMallocHook(size_t, const void *);
static void customFreeHook(void*, const void*);

// Static variable to backup standard malloc hook
static void *(*systemMalloc)(size_t, const void *);
static void (*systemFree)(void*, const void *);

// Static variable overrides malloc at application startup time
//void (*__malloc_initialize_hook) (void) = installMemoryWatcher;

typedef struct _pageEntry
{
	void *address;
	int size;
} pageEntry;

static int pageSize = -1;
static vector<pageEntry> allocatedPages;
static vector<void*> dirtyPages;

static void *customMallocHook(size_t size, const void *caller)
{
	void *result;

	// Restore standard malloc
	__malloc_hook = systemMalloc;

	size_t totalSize = ( ( size + (pageSize-1) ) / pageSize ) * pageSize;

	systemMalloc = __malloc_hook;

	// Call standard malloc
	//   result = malloc(size);
	result = memalign(pageSize, totalSize);

	pageEntry newEntry;
	newEntry.address = result;
	newEntry.size = totalSize;

	allocatedPages.push_back(newEntry);

//	printf("\nTrapped malloc(%u)\n", (unsigned int) size);
//	printf("\tAllocated 0x%x bytes\n", totalSize);
//	printf("\tAdded %p to allocated pages\n", result);

	// Reinstall custom malloc
	__malloc_hook = customMallocHook;

	return result;
}

static void customFreeHook(void *ptr, const void *caller)
{
	// Restore standard free
	__free_hook = systemFree;

	// Call standard malloc
	free(ptr);

//	allocatedPages.find

	for( size_t i = 0 ; i < allocatedPages.size(); ++i )
	{
		if( allocatedPages[i].address == ptr )
		{
			allocatedPages.erase( allocatedPages.begin() + i );
//			printf("\nRemoved area(%p)\n", ptr);
			break;
		}
	}

	systemFree = __free_hook;

//	printf("\nTrapped free(%p)\n", ptr);

	// Reinstall custom free
	__free_hook = customFreeHook;
}

static void segvHandler(int sig, siginfo_t *si, void *unused)
{
	static void *(*tempMalloc)(size_t, const void *);
	static void (*tempFree)(void*, const void *);

	tempMalloc = __malloc_hook;
	tempFree = __free_hook;

	// Restore standard malloc
	__malloc_hook = systemMalloc;

	// Restore standard free
	__free_hook = systemFree;

//	printf("\nTrapped SIGSEGV at address: %p\n", si->si_addr);

	unsigned long address = (unsigned long)si->si_addr & ~(pageSize-1);

//	printf("\tUnprotecting page at address: %p\n", (void*)address);

	if( mprotect((void*)(address), pageSize, PROT_READ | PROT_WRITE) == -1 )
	{
		handle_error("mprotect");
	}

//	printf("\tAdding page at address: %p to dirty pages array\n", (void*)address);

	// Push back address of dirty page
	dirtyPages.push_back((void*)(address));

	// Reinstall custom malloc
	__malloc_hook = tempMalloc;

	// Reinstall custom free
	__free_hook = tempFree;

}

void installMemoryWatcher()
{
	allocatedPages.reserve(16384);
	dirtyPages.reserve(16384);

	pageSize = sysconf(_SC_PAGE_SIZE);

	if (pageSize == -1)
	{
		handle_error("sysconf");
	}

	systemMalloc = __malloc_hook; // Save original malloc
	__malloc_hook = customMallocHook; // Install custom malloc

	systemFree = __free_hook;
	__free_hook = customFreeHook;

//	printf("\nMemory watcher installed\n");
//	printf("\tPage size = 0x%x\n", pageSize);
}

void removeMemoryWatcher()
{
	__malloc_hook = systemMalloc;
	__free_hook = systemFree;

//	printf("\nMemory watcher removed\n\n");
}

void protectAll()
{
	struct sigaction sa;

//	printf("\nInstalling segv handler\n");

	sa.sa_flags = SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	sa.sa_sigaction = segvHandler;
	if( sigaction(SIGSEGV, &sa, NULL) == -1 )
	{
		handle_error("sigaction");
	}

//	printf("Protecting allocated pages\n");

	for( unsigned int page = 0 ; page < allocatedPages.size(); ++page )
	{
		void *address = allocatedPages[page].address;

//		printf("\tProtecting %p\n", address);

		if( mprotect(address, allocatedPages[page].size, PROT_READ) == -1 )
		{
			handle_error("mprotect");
		}
	}
}

void unprotectAll()
{
	struct sigaction sa;

//	printf("\nUnprotecting allocated pages\n");

	for( unsigned int page = 0 ; page < allocatedPages.size(); ++page )
	{
		void *address = allocatedPages[page].address;

//		printf("\tUnprotecting %p\n", address);

		if( mprotect(address, pageSize, PROT_READ | PROT_WRITE) == -1 )
		{
			handle_error("mprotect");
		}
	}

//	printf("Removing segv handler\n");

	sa.sa_flags = 0;
	sigemptyset(&sa.sa_mask);
	sa.sa_handler = SIG_DFL;

	if( sigaction(SIGSEGV, &sa, NULL) == -1 )
	{
		handle_error("sigaction");
	}
}

void listDirtyPages()
{
//	printf("\nDirty pages list\n");

	for( unsigned int page = 0 ; page < dirtyPages.size(); ++page )
	{
//		printf("\tPage:%p\n", dirtyPages[page]);
	}
}

const vector<void*>& getDirtyPages()
{
	return dirtyPages;
}

unsigned int getPageSize()
{
	return sysconf(_SC_PAGE_SIZE);
}
